# This code is part of LINKEQ.
#
# (C) Copyright LINKE 2023.
#
# This code is licensed under the Apache License, Version 2.0. You may
# obtain a copy of this license in the LICENSE.txt file in the root directory
# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
#
# Any modifications or derivative works of this code must retain this
# copyright notice, and modified files need to carry a notice indicating
# that they have been altered from the originals.
#
# -*- coding: utf-8 -*-
# @Time     : 2023/6/7 10:48
# @Author   : HFALSH @ LINKE
# @File     : fake_backend.py
# @Software : PyCharm
import numpy as np
from typing import Union, List, Tuple, Callable, Dict, Any, Optional, Iterator, Iterable, TypeVar
from quantumcircuit.gate import Gate
import scipy.stats as stats
from quantumcircuit.circuit import QuantumCircuit
from backend import Backend


class FakeChip(Backend):
    def __init__(self,
                 name: str,
                 n_qubits: int,
                 gate_duration: Optional[dict] = None,
                 coupling_map: Optional[List[Tuple[int, int]]] = None,
                 basis_gates: Optional[List[str]] = None,
                 qubit_t1: Optional[Union[float, List[float]]] = None,
                 qubit_t2: Optional[Union[float, List[float]]] = None,
                 single_qubit_noise: [Union[float, Dict[str, float]]] = None,
                 two_qubit_noise: Optional[Union[float, Dict[str, float]]] = None,
                 std: Optional[float] = None,
                 seed: Optional[int] = None
                 ):
        """
        Initialize a fake chip.
        Args:
            name:
            n_qubits:
            gate_duration: the execution time of each gate
            coupling_map: chip topology
            basis_gates: gate set supported by the chip
            qubit_t1: t1 coherence time of each qubit on the chip; input a list or expectation
            qubit_t2: t2 coherence time of each qubit on the chip; input a list or expectation
            single_qubit_noise: single-qubit gate noise of each qubit of the chip, suppose all the
                single-qubit gate noise are the same
            two_qubit_noise: two-qubit gate noise of each connected edge on the topology graph，only cx gate now
            std:
            seed:
        """
        np.random.seed(seed)

        if basis_gates is None:
            basis_gates = ["X", "Y", "Z", "H", "S", "T", "CX", "RX", "RZ"]

        if std is None:
            std = 0.01

        if not isinstance(coupling_map, list):
            coupling_map = [(i, j) for j in range(n_qubits) for i in range(n_qubits)]

        if not isinstance(qubit_t1, list):
            qubit_t1 = np.random.normal(loc=qubit_t1 or 113.0, scale=std * 100, size=n_qubits).tolist()

        if not isinstance(qubit_t2, list):
            qubit_t2 = np.random.normal(loc=qubit_t2 or 150.2, scale=std * 100, size=n_qubits).tolist()

        if not isinstance(gate_duration, list):
            gate_duration = {}
            single_qubit_set = ["X", "Y", "Z", "H", "S", "T", "RX", "RZ"]
            double_qubit_set = ["CX"]
            # qubit [(name, [],time)]
            for qubit in range(n_qubits):
                times = np.random.normal(loc=200, scale=50, size=len(single_qubit_set)).tolist()
                for i, gate in enumerate(single_qubit_set):
                    key = str(qubit) + ',' + gate + ',' + str(qubit)
                    gate_duration[key] = times[i]
            for qubit in range(n_qubits):
                involved_two_qubit_gates = []
                num_involved_gates = 0
                for coupling in coupling_map:
                    if qubit == coupling[0]:
                        num_involved_gates += 1
                        involved_two_qubit_gates.append(coupling)
                two_qubit_time = np.random.normal(loc=600, scale=200,
                                                  size=num_involved_gates * len(double_qubit_set)).tolist()
                for i, gate in enumerate(double_qubit_set):
                    for j, qubit_pair in enumerate(involved_two_qubit_gates):
                        key = str(qubit) + ',' + gate + ',' + str(qubit_pair)
                        gate_duration[key] = two_qubit_time[i * len(double_qubit_set) + j]
            # print(gate_duration)

        if not isinstance(single_qubit_noise, list):
            lower = 0
            upper = 1
            expectation = single_qubit_noise or 0.99
            X = stats.truncnorm((lower - expectation) / std, (upper - expectation) / std, loc=expectation, scale=std)
            single_qubit_noise_list = X.rvs(n_qubits)
            single_qubit_noise = {}
            for qubit in range(n_qubits):
                single_qubit_noise[str(qubit)] = single_qubit_noise_list[qubit]
            # single_qubit_noise = np.random.normal(loc=single_qubit_noise or 0.99, scale=std, size=n_qubits).tolist()

        if not isinstance(two_qubit_noise, list):
            lower = 0
            upper = 1
            expectation = two_qubit_noise or 0.95
            X = stats.truncnorm((lower - expectation) / std, (upper - expectation) / std, loc=expectation, scale=std)
            two_qubit_noise_list = X.rvs(len(coupling_map))
            two_qubit_noise = {}
            for i, coupling in enumerate(coupling_map):
                coupling_key = str(coupling[0]) + '_' + str(coupling[1])
                two_qubit_noise[coupling_key] = two_qubit_noise_list[i]

        super().__init__(name)
        self.qubit_t2 = qubit_t2
        self.qubit_t1 = qubit_t1
        self.n_qubits = n_qubits
        self.basis_gates = basis_gates
        self.coupling_map = coupling_map
        self.name = name
        self.gate_duration = gate_duration
        self.single_qubit_noise = single_qubit_noise
        self.two_qubit_noise = two_qubit_noise
        # qubit [(name, [],time)]

    def get_backend_info(self):
        return f"Backend: {self.name}"

    def set_all_noise(self, noise_dict: Dict):
        """
        Args:
            noise_dict: the dictionary of noise on a quantum device
            noise_dict = {'0': 0.0012, '1': 0.002, '0_1': 0.015}
            IBM_guadalupe_noise_dict = {'0': 0.0002112, '1': 0.0004517, '2': 0.000415, '3': 0.000238,
            '4': 0.000256, '5': 0.0003913, '6': 0.000445, '7': 0.0002947, '8': 0.0003818, '9': 0.0001998,
             '10': 0.0002107, '11': 0.0006221, '12': 0.000249, '13': 0.0004251, '14': 0.0002433,
             '15': 0.0001669, '0_1': 0.01055, '1_0': 0.01055, '1_4': 0.007236, '4_1': 0.007236,
             '1_2': 0.009543, '2_1': 0.009543, '2_3': 0.01606, '3_2': 0.01606, '3_5': 0.01091,
             '5_3': 0.01091, '4_7': 0.01336, '7_4': 0.01336, '5_8': 0.008527, '8_5': 0.008527,
             '6_7': 0.0112, '7_6': 0.0112, '7_10': 0.006437, '10_7': 0.006437, '8_11': 0.02012,
             '11_8': 0.02012, '8_9': 0.01428, '9_8': 0.01428, '10_12': 0.01426, '12_10': 0.01426,
             '11_14': 0.01728, '14_11': 0.01728, '12_15': 0.007435, '15_12': 0.007435, '12_13': 0.006903,
             '13_12': 0.006903, '13_14': 0.01054, '14_13': 0.01054}

        Returns:
            Setting the noise of single qubit gate and two-qubit gate noise as the input dict describes.
        """
        self.two_qubit_noise = {}
        self.single_qubit_noise = {}
        for key in noise_dict.keys():
            if '_' not in key:
                assert int(key) < self.n_qubits
                self.single_qubit_noise[key] = noise_dict[key]
            elif '_' in key:
                involved_qubits = key.split('_')
                assert int(involved_qubits[0]) < self.n_qubits
                assert int(involved_qubits[1]) < self.n_qubits
                assert (int(involved_qubits[0]), int(involved_qubits[1])) in self.coupling_map or (
                int(involved_qubits[1]), int(involved_qubits[0])) in self.coupling_map
                self.two_qubit_noise[key] = noise_dict[key]

    def run(self, qcircuit, *args):
        assert type(qcircuit) == QuantumCircuit
        pass


if __name__ == "__main__":
    fake_test = FakeChip(name="Faker", n_qubits=16)
    # print(fake_test.gate_duration)
    # print(fake_test.qubit_t1)
    # print(fake_test.qubit_t2)
    IBM_guadalupe_noise_dict = {'0': 0.0002112, '1': 0.0004517, '2': 0.000415, '3': 0.000238,
                                '4': 0.000256, '5': 0.0003913, '6': 0.000445, '7': 0.0002947, '8': 0.0003818,
                                '9': 0.0001998,
                                '10': 0.0002107, '11': 0.0006221, '12': 0.000249, '13': 0.0004251, '14': 0.0002433,
                                '15': 0.0001669, '0_1': 0.01055, '1_0': 0.01055, '1_4': 0.007236, '4_1': 0.007236,
                                '1_2': 0.009543, '2_1': 0.009543, '2_3': 0.01606, '3_2': 0.01606, '3_5': 0.01091,
                                '5_3': 0.01091, '4_7': 0.01336, '7_4': 0.01336, '5_8': 0.008527, '8_5': 0.008527,
                                '6_7': 0.0112, '7_6': 0.0112, '7_10': 0.006437, '10_7': 0.006437, '8_11': 0.02012,
                                '11_8': 0.02012, '8_9': 0.01428, '9_8': 0.01428, '10_12': 0.01426, '12_10': 0.01426,
                                '11_14': 0.01728, '14_11': 0.01728, '12_15': 0.007435, '15_12': 0.007435,
                                '12_13': 0.006903,
                                '13_12': 0.006903, '13_14': 0.01054, '14_13': 0.01054}
    fake_test.set_all_noise(IBM_guadalupe_noise_dict)
    print(fake_test.coupling_map)
    print(fake_test.single_qubit_noise)
    print(fake_test.two_qubit_noise)

    # lower, upper = 0, 1
    # mu, sigma = 0.99, 0.01
    # X = stats.truncnorm(
    #     (lower - mu) / sigma, (upper - mu) / sigma, loc=mu, scale=sigma)
    # N = stats.norm(loc=mu, scale=sigma)
    # print(X.rvs(10))
    # print((max(X.rvs(10000)) > 1))
    # print((max(N.rvs(10000)) > 1))
